#version 450

#define ambient 0.1

layout (input_attachment_index = 0, binding = 0) uniform subpassInput inputPosition;
layout (input_attachment_index = 1, binding = 1) uniform subpassInput inputNormal;
layout (input_attachment_index = 2, binding = 2) uniform subpassInput inputAlbedo;

layout (push_constant) uniform PushConstants {
    int attachmentIndex;
} pc;

struct PointLight {
    vec4 position;
    vec3 color;
    float radius;
};
layout(std140, binding = 3) buffer LightsBuffer {
    PointLight lights[];
};

layout(location=0) out vec4 fragColor;

vec3 calLight(vec3 L, vec3 N, vec3 lightColor) {
    return max(0.0, dot(N, L)) * lightColor;
}

void main() {
    if (pc.attachmentIndex == 0) {
        fragColor.rgb = subpassLoad(inputPosition).rgb;
    } else if (pc.attachmentIndex == 1) {
        fragColor.rgb = subpassLoad(inputNormal).rgb;
    } else if (pc.attachmentIndex == 2) {
        fragColor = subpassLoad(inputAlbedo);
    } else {
        vec3 fragPos = subpassLoad(inputPosition).rgb;
        vec3 normal = subpassLoad(inputNormal).rgb;
        vec4 albedo = subpassLoad(inputAlbedo);

        vec3 color = albedo.rgb * ambient;
        vec3 N = normalize(normal);

        for(int i = 0; i < lights.length(); ++i) {
            vec3 L = lights[i].position.xyz - fragPos;
            float dist = length(L);
            float atten = lights[i].radius / (pow(dist, 3.0) + 1.0);
            color += calLight(normalize(L), N, lights[i].color) * albedo.rgb * atten;
        }
        fragColor = vec4(color, albedo.a);
    }
}